#include <opencv2/core/core.hpp>
#include <opencv2/calib3d/calib3d.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <stdio.h>
#include <iostream>
#include "popt_pp.h"
#include <sys/stat.h>

using namespace std;
using namespace cv;

/*存棋盘内角点三维坐标*/
vector<vector<Point3f> > object_points;
/*存棋盘内角点二维坐标*/
vector<vector<Point2f> > image_points;


Mat img, gray;
Size im_size;

bool doesExist(const std::string &name) {
    struct stat buffer;
    return (stat(name.c_str(), &buffer) == 0);
}

void visulise_chess(Mat gray, Size board_size, bool found, String name, vector<Point2f> corners, int wait_time = 500) {
    Mat img_tmp;
    gray.copyTo(img_tmp);
    drawChessboardCorners(img_tmp, board_size, corners, found);
    imshow(name, img_tmp);//显示图片
    waitKey(wait_time);
}

void setup_calibration(int board_width, int board_height, int num_imgs,
                       float square_size, char *imgs_directory, char *imgs_filename,
                       char *extension) {
    vector<Point2f> corners;
    Size board_size = Size(board_width, board_height);
    int board_n = board_width * board_height;

    for (int k = 1; k <= num_imgs; k++) {
        char img_file[100];
        sprintf(img_file, "%s%s%d.%s", imgs_directory, imgs_filename, k, extension);
        if (!doesExist(img_file))
            continue;
        img = imread(img_file, CV_LOAD_IMAGE_COLOR);
        cv::cvtColor(img, gray, CV_BGR2GRAY);

        bool found = false;
        /* 提取棋盘角点 */
        found = cv::findChessboardCorners(img, board_size, corners,
                                          CV_CALIB_CB_ADAPTIVE_THRESH | CV_CALIB_CB_FILTER_QUADS);
        if (found) {
//            visulise_chess(gray, board_size, found, "Camera Calibration",corners, 500);
            /* 亚像素精确化 */
            cornerSubPix(gray, corners, cv::Size(5, 5), cv::Size(-1, -1),
                         TermCriteria(CV_TERMCRIT_EPS | CV_TERMCRIT_ITER, 30, 0.1));

//            visulise_chess(gray, board_size, found, "Camera Calibration cornerSubPix",corners, 500);
        }

        /*计算相对于棋盘角点原点的三维坐标（无深度），需要棋盘方块边长，这里单位是m*/
        vector<Point3f> obj;
        for (int i = 0; i < board_height; i++)
            for (int j = 0; j < board_width; j++)
                obj.push_back(Point3f((float) j * square_size, (float) i * square_size, 0));

        if (found) {
            /*三维坐标坐标原点是 棋盘内角点原点，二维坐标原点是图像左上角原点*/
            cout << k << ". Found corners!" << endl;
            image_points.push_back(corners);
            object_points.push_back(obj);
        }
    }
}

double computeReprojectionErrors(const vector<vector<Point3f> > &objectPoints,
                                 const vector<vector<Point2f> > &imagePoints,
                                 const vector<Mat> &rvecs, const vector<Mat> &tvecs,
                                 const Mat &cameraMatrix, const Mat &distCoeffs) {
    vector<Point2f> imagePoints2;
    int i, totalPoints = 0;
    double totalErr = 0, err;
    vector<float> perViewErrors;
    perViewErrors.resize(objectPoints.size());

    for (i = 0; i < (int) objectPoints.size(); ++i) {
        projectPoints(Mat(objectPoints[i]), rvecs[i], tvecs[i], cameraMatrix,
                      distCoeffs, imagePoints2);
        err = norm(Mat(imagePoints[i]), Mat(imagePoints2), CV_L2);
        int n = (int) objectPoints[i].size();
        perViewErrors[i] = (float) std::sqrt(err * err / n);
        totalErr += err * err;
        totalPoints += n;
    }
    return std::sqrt(totalErr / totalPoints);
}

/* cmd: ./calibrate -w 9 -h 6 -n 27 -s 0.02423 -d "../calib_imgs/1/" -i "left" -o "cam_left.yml" -e "jpg" */
/* ******************************************************************************
 * -w 棋盘横向内角点数
 * -h 棋盘纵向内角点数
 * -n 图片数量
 * -s 棋盘方块边长（单位m）
 * -d 目录
 * -i 文件名
 * -o 输出
 * -e 扩展名
 * *******************************************************************************/
int main(int argc, char const **argv) {
    int board_width, board_height, num_imgs;
    float square_size;
    char *imgs_directory;
    char *imgs_filename;
    char *out_file;
    char *extension;

    static struct poptOption options[] = {
            {"board_width", 'w', POPT_ARG_INT, &board_width, 0, "Checkerboard width", "NUM"},
            {"board_height", 'h', POPT_ARG_INT, &board_height, 0, "Checkerboard height", "NUM"},
            {"num_imgs", 'n', POPT_ARG_INT, &num_imgs, 0, "Number of checkerboard images", "NUM"},
            {"square_size", 's', POPT_ARG_FLOAT, &square_size, 0, "Size of checkerboard square", "NUM"},
            {"imgs_directory", 'd', POPT_ARG_STRING, &imgs_directory, 0, "Directory containing images", "STR"},
            {"imgs_filename", 'i', POPT_ARG_STRING, &imgs_filename, 0, "Image filename", "STR"},
            {"extension", 'e', POPT_ARG_STRING, &extension, 0, "Image extension", "STR"},
            {"out_file", 'o', POPT_ARG_STRING, &out_file, 0, "Output calibration filename (YML)", "STR"},
            POPT_AUTOHELP
            {NULL, 0, 0, NULL, 0, NULL, NULL}
    };

    POpt popt(NULL, argc, argv, options, 0);
    int c;
    while ((c = popt.getNextOpt()) >= 0) {}

    setup_calibration(board_width, board_height, num_imgs, square_size,
                      imgs_directory, imgs_filename, extension);

    printf("Starting Calibration\n");
    Mat K;
    Mat D;
    vector<Mat> rvecs, tvecs;
    int flag = 0;
    flag |= CV_CALIB_FIX_K4;
    flag |= CV_CALIB_FIX_K5;
    calibrateCamera(object_points, image_points, img.size(), K, D, rvecs, tvecs, flag);

    cout << "Calibration error: " << computeReprojectionErrors(object_points, image_points, rvecs, tvecs, K, D) << endl;

    FileStorage fs(out_file, FileStorage::WRITE);
    fs << "K" << K;
    fs << "D" << D;
    fs << "board_width" << board_width;
    fs << "board_height" << board_height;
    fs << "square_size" << square_size;
    printf("Done Calibration\n");

    cout << "K : " << K << endl;
    cout << "D : " << D << endl;
    cout << "board_width : " << board_width << endl;
    cout << "board_height : " << board_height << endl;
    cout << "square_size : " << square_size << endl;

    return 0;
}
